{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### And welcome to Week 4, Day 3 - more LangGraph.."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Importing necessary libraries\n",
    "\n",
    "import requests\n",
    "import os\n",
    "import gradio as gr\n",
    "from typing import TypedDict\n",
    "from dotenv import load_dotenv\n",
    "from IPython.display import Image, display\n",
    "from typing import Annotated\n",
    "\n",
    "# Using Amazon Bedrock Models\n",
    "from langchain_aws import ChatBedrockConverse \n",
    "from langgraph.graph import StateGraph, START, END\n",
    "from langgraph.graph.message import add_messages\n",
    "from langgraph.graph import StateGraph\n",
    "from langgraph.prebuilt import ToolNode, tools_condition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loading up environment variables\n",
    "# Using amazon bedrock, values like aws secret key, access key, bedrock api key and aws region name should exist\n",
    "load_dotenv(override=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set up LangSmith via the link below to enable tracing:\n",
    "\n",
    "https://langsmith.com"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Next, here is a useful function in LangChain community:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.utilities import GoogleSerperAPIWrapper\n",
    "\n",
    "serper = GoogleSerperAPIWrapper()\n",
    "serper.run(\"What is the capital of Kenya?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### LangChain wrapper class for converting functions into Tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import Tool\n",
    "\n",
    "search_tool = Tool(\n",
    "    name=\"Internet_Search_Tool\",\n",
    "    func=serper.run,\n",
    "    description=\"Used to get information from the internet.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Invoking the tool in the langchain way"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "search_tool.invoke(\"What is the capital of Kenya?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Push Notification Tool\n",
    "\n",
    "This is a push notification tool, that when invoked it sends an sms push notification to the phone."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pushover_token = os.getenv(\"PUSHOVER_TOKEN\")\n",
    "pushover_user = os.getenv(\"PUSHOVER_USER\")\n",
    "pushover_url = \"https://api.pushover.net/1/messages.json\"\n",
    "\n",
    "def push(text: str):\n",
    "    \"\"\"Send a push notification to the user\"\"\"\n",
    "    requests.post(pushover_url, \n",
    "                    data = {\"token\": pushover_token, \n",
    "                            \"user\": pushover_user, \n",
    "                            \"message\": text})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "push_notification_tool = Tool(\n",
    "    name=\"Send_Push_Notification\",\n",
    "    func=push,\n",
    "    description=\"Used to send a push notification to the phone.\"\n",
    ")\n",
    "\n",
    "#push_notification_tool.invoke(\"Hello. This is a test notication\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using TypedDict instead of BaseModel for the State object\n",
    "\n",
    "When we implement tools, we always need to make 2 changes to the code:\n",
    "\n",
    "1. Changes to provide the tools to Amazon Bedrock in json when we make the call\n",
    "\n",
    "2. Changes to handle the results back: look for the model staying that the finish_reason==\"tool_calls\" and then retrieve the call, run the function, provide the results."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Bring the tools together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tools = [search_tool, push_notification_tool]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 1: Define the State object\n",
    "class State(TypedDict):\n",
    "    messages: Annotated[list, add_messages]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 2: Start the Graph Builder with this State class\n",
    "graph_builder = StateGraph(State)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Invoking the chat bedrock api using amazon nova micro model\n",
    "\n",
    "llm = ChatBedrockConverse(model_id=\"amazon.nova-micro-v1:0\")\n",
    "\n",
    "# Binding the LLMs with tools\n",
    "llm_with_tools = llm.bind_tools(tools)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 3: Create a Node\n",
    "from langgraph import graph\n",
    "\n",
    "def chatbot(state: State):\n",
    "    return {\"messages\": [llm_with_tools.invoke(state[\"messages\"])]}\n",
    "\n",
    "graph_builder.add_node(\"chatbot\", chatbot)\n",
    "graph_builder.add_node(\"tools\", ToolNode(tools=tools))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 4: Create Edges\n",
    "\n",
    "graph_builder.add_conditional_edges(\"chatbot\", tools_condition, \"tools\")\n",
    "\n",
    "# Any time a tool is called, we return to the chatbot to decide the next step\n",
    "graph_builder.add_edge(\"tools\", \"chatbot\")\n",
    "graph_builder.add_edge(START, \"chatbot\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 5: Compile the Graph and display the flow\n",
    "graph = graph_builder.compile()\n",
    "display(Image(graph.get_graph().draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Creating the chatbot function and displating gradio intereface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def chat(user_input: str, history):\n",
    "    result = graph.invoke({\"messages\": [{\"role\": \"user\", \"content\": user_input}]})\n",
    "    return result[\"messages\"][-1].content\n",
    "\n",
    "\n",
    "gr.ChatInterface(chat, type=\"messages\").launch()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adding Memory\n",
    "\n",
    "We have this whole Graph maintaining the state and appending to the state.\n",
    "\n",
    "Why isn't this handling memory?\n",
    "\n",
    "This is a crucial point for understanding LangGraph\n",
    "\n",
    "> A super-step can be considered a single iteration over the graph nodes. Nodes that run in parallel are part of the same super-step, while nodes that run sequentially belong to separate super-steps.\n",
    "\n",
    "\n",
    "One \"Super-Step\" of the graph represents one invocation of passing messages between agents.\n",
    "\n",
    "In idomatic LangGraph, you call invoke to run your graph for each super-step; for each interaction.\n",
    "\n",
    "The `reducer` handles state updates automatically within one super-step, but not between them.\n",
    "\n",
    "That is what `checkpointing` achieves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "\n",
    "memory = MemorySaver()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Steps 1 and 2\n",
    "graph_builder = StateGraph(State)\n",
    "\n",
    "# Step 3\n",
    "llm = ChatBedrockConverse(model_id=\"amazon.nova-micro-v1:0\")\n",
    "llm_with_tools = llm.bind_tools(tools)\n",
    "\n",
    "def chatbot(state: State):\n",
    "    print(state)\n",
    "    return {\"messages\": [llm_with_tools.invoke(state[\"messages\"])]}\n",
    "\n",
    "graph_builder.add_node(\"chatbot\", chatbot)\n",
    "graph_builder.add_node(\"tools\", ToolNode(tools=tools))\n",
    "\n",
    "# Step 4\n",
    "graph_builder.add_conditional_edges( \"chatbot\", tools_condition, \"tools\")\n",
    "graph_builder.add_edge(\"tools\", \"chatbot\")\n",
    "graph_builder.add_edge(START, \"chatbot\")\n",
    "\n",
    "# Step 5\n",
    "graph = graph_builder.compile(checkpointer=memory)\n",
    "display(Image(graph.get_graph().draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pick a thread to use as the key for this conversation\n",
    "config = {\"configurable\": {\"thread_id\": \"1\"}}\n",
    "\n",
    "# Chatbot config\n",
    "def chat(user_input: str, history):\n",
    "    results = graph.invoke({\"messages\": [{\"role\": \"user\", \"content\": user_input}]}, config=config)\n",
    "    return results[\"messages\"][-1].content\n",
    "\n",
    "# # launching gradio\n",
    "gr.ChatInterface(chat, type=\"messages\").launch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# getting the state of the config\n",
    "graph.get_state(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Listing the state of the config, Most recent first\n",
    "\n",
    "list(graph.get_state_history(config))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### LangGraph gives you tools to set the state back to a prior point in time, to branch off:\n",
    "\n",
    "```\n",
    "config = {\"configurable\": {\"thread_id\": \"1\", \"checkpoint_id\": ...}}\n",
    "graph.invoke(None, config=config)\n",
    "```\n",
    "\n",
    "And this allows you to build stable systems that can be recovered and rerun from any prior checkpoint."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " #### Memory - SQL\n",
    " And now let's store in SQL\n",
    "\n",
    " And this is the power of LangGraph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# importing and setting up sqlite db\n",
    "\n",
    "import sqlite3\n",
    "from langgraph.checkpoint.sqlite import SqliteSaver\n",
    "\n",
    "db_path = \"memory.db\"\n",
    "connection = sqlite3.connect(db_path, check_same_thread=False)\n",
    "sql_memory = SqliteSaver(connection)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Steps 1 and 2\n",
    "graph_builder = StateGraph(State)\n",
    "\n",
    "# Step 3\n",
    "llm = ChatBedrockConverse(model_id=\"amazon.nova-micro-v1:0\")\n",
    "tools = [search_tool, push_notification_tool]\n",
    "llm_with_tools = llm.bind_tools(tools)\n",
    "\n",
    "def chatbot(state: State):\n",
    "    print(state)\n",
    "    return {\"messages\": [llm_with_tools.invoke(state[\"messages\"])]}\n",
    "\n",
    "graph_builder.add_node(\"chatbot\", chatbot)\n",
    "graph_builder.add_node(\"tools\", ToolNode(tools=tools))\n",
    "\n",
    "# Step 4\n",
    "graph_builder.add_conditional_edges( \"chatbot\", tools_condition, \"tools\")\n",
    "graph_builder.add_edge(\"tools\", \"chatbot\")\n",
    "graph_builder.add_edge(START, \"chatbot\")\n",
    "\n",
    "# Step 5\n",
    "graph = graph_builder.compile(checkpointer=sql_memory)\n",
    "display(Image(graph.get_graph().draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\"configurable\": {\"thread_id\": \"2\"}}\n",
    "\n",
    "def chat(user_input: str, history):\n",
    "    result = graph.invoke({\"messages\": [{\"role\": \"user\", \"content\": user_input}]}, config=config)\n",
    "    return result[\"messages\"][-1].content\n",
    "\n",
    "\n",
    "gr.ChatInterface(chat, type=\"messages\").launch()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
